import os
import re
import matplotlib.pyplot as plt

# Define regex patterns for parsing the data
energy_pattern = re.compile(r'FINAL_ETOT_IS ([\-0-9\.]+) eV')
magnetism_pattern = re.compile(r'Total Magnetism on atom:\s+Fe\s+([0-9\.]+)')

# Initialize data storage
energies = {}
magnetisms = {}

# Read energy values from running_scf.log files
for dirname in os.listdir('.'):
    if os.path.isdir(dirname):
        energy_file = os.path.join(dirname, 'OUT.bak', 'running_scf.log')
        if os.path.isfile(energy_file):
            with open(energy_file, 'r') as file:
                for line in file:
                    match = energy_pattern.search(line)
                    if match:
                        energies[dirname] = float(match.group(1))
                        break

# Read magnetism values from mulliken.txt files
for dirname in os.listdir('.'):
    if os.path.isdir(dirname):
        mag_file = os.path.join(dirname, 'OUT.bak', 'mulliken.txt')
        if os.path.isfile(mag_file):
            with open(mag_file, 'r') as file:
                for line in file:
                    match = magnetism_pattern.search(line)
                    if match:
                        magnetisms[dirname] = float(match.group(1))
                        break

# Calculate relative energies with respect to the ground state
ground_state_energy = energies.get('a7')
relative_energies = {k: (v - ground_state_energy)*1000 for k, v in energies.items()}

# Prepare data for plotting
#magnetism_values = [magnetisms[k] for k in sorted(magnetisms.keys())]
magnetism_values = [4.0,5.0,6.0,7.0,8.0,9.0,10.0]
relative_energy_values = [relative_energies[k] for k in sorted(relative_energies.keys())]
labels = sorted(magnetisms.keys())

# Plotting the data
plt.figure(figsize=(10, 5))
plt.plot(magnetism_values, relative_energy_values, 'ko-')

# Highlight the dgs data point with a red circle
if 'gs' in magnetisms and 'gs' in relative_energies:
    plt.plot(magnetisms['gs'], relative_energies['gs'], 'ro', label='Ground state')

plt.xlabel('Angle (deg)')
plt.ylabel('Relative Energy (meV)')
plt.title('Noncollinear Fe$_3$ ground state')
#plt.legend()
plt.grid(True)

# Save the plot as a PNG file
plt.savefig('energy_vs_angle.png')
print("Plot saved as 'energy_vs_angle.png'.")
plt.show()
